package com.tancire.common.utils;

import com.tancire.common.exception.base.RsaException;
import com.tancire.common.utils.bean.PairBean;
import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

public class RsaUtils {

    /**
     * 加密方式
     */
    public static final String CIPHER_KEY = "RSA";

    /**
     * 默认
     */
    public static final int DEFAULT_KEY_SIZE = 1024;

    /**
     * 加密
     */
    public static final int ENCRYPT = 1;

    /**
     * 解密
     */
    public static final int DECRYPT = 2;

    /**
     * 生成RSA密钥对
     *
     * @return
     */
    public static PairBean generateRsaPair() {
        return generateRsaPair(DEFAULT_KEY_SIZE);
    }


    /**
     * 生成RSA密钥对
     *
     * @param keySize
     * @return
     */
    public static PairBean generateRsaPair(int keySize) {
        PairBean resPair = new PairBean();
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(CIPHER_KEY);
            keyPairGenerator.initialize(keySize);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            RSAPublicKey rsaPublicKey = (RSAPublicKey) keyPair.getPublic();
            RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) keyPair.getPrivate();
            resPair.setPublicKey(Base64.encodeBase64String(rsaPublicKey.getEncoded()));
            resPair.setPrivateKey(Base64.encodeBase64String(rsaPrivateKey.getEncoded()));
        } catch (Exception e) {
            throw new RsaException("创建密钥失败");
        }
        return resPair;
    }


    /**
     * 公钥
     * 加密/解密
     *
     * @return
     */
    public static String pubRsa(String letter, String publicKey, int mode) {
        try {
            byte[] bytes = Base64.decodeBase64(publicKey);
            X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(bytes);
            KeyFactory keyFactory = KeyFactory.getInstance(CIPHER_KEY);
            PublicKey key = keyFactory.generatePublic(x509EncodedKeySpec);
            Cipher cipher = Cipher.getInstance(CIPHER_KEY);
            cipher.init(mode, key);
            return codeHandle(letter, cipher, mode);
        } catch (Exception e) {
            if (Cipher.ENCRYPT_MODE == mode) {
                throw new RsaException("加密数据失败");
            } else {
                throw new RsaException("解密数据失败");
            }
        }
    }

    /**
     * 公钥
     * 加密/解密
     *
     * @return
     */
    public static String priRsa(String letter, String privateKey, int mode) {
        try {
            byte[] bytes = Base64.decodeBase64(privateKey);
            PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(bytes);
            KeyFactory keyFactory = KeyFactory.getInstance(CIPHER_KEY);
            PrivateKey key = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
            Cipher cipher = Cipher.getInstance(CIPHER_KEY);
            cipher.init(mode, key);
            return codeHandle(letter, cipher, mode);
        } catch (Exception e) {
            if (Cipher.ENCRYPT_MODE == mode) {
                throw new RsaException("加密数据失败");
            } else {
                throw new RsaException("解密数据失败");
            }
        }
    }

    /**
     * 加密
     * 公钥加密不需要传私钥，私钥加密不需要传公钥
     *
     * @return
     */
    public static String rsaEncrypt(String publicKey, String privateKey, String letter) {
        try {
            if (null != publicKey && null == privateKey) {
                byte[] bytes = Base64.decodeBase64(publicKey);
                X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(bytes);
                KeyFactory keyFactory = KeyFactory.getInstance(CIPHER_KEY);
                PublicKey key = keyFactory.generatePublic(x509EncodedKeySpec);
                Cipher cipher = Cipher.getInstance(CIPHER_KEY);
                cipher.init(Cipher.ENCRYPT_MODE, key);
                byte[] res = cipher.doFinal(letter.getBytes(StandardCharsets.UTF_8));
                return Base64.encodeBase64String(res);
            } else if (null == publicKey && null != privateKey) {
                byte[] bytes = Base64.decodeBase64(privateKey);
                PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(bytes);
                KeyFactory keyFactory = KeyFactory.getInstance(CIPHER_KEY);
                PrivateKey key = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
                Cipher cipher = Cipher.getInstance(CIPHER_KEY);
                cipher.init(Cipher.ENCRYPT_MODE, key);
                byte[] res = cipher.doFinal(letter.getBytes(StandardCharsets.UTF_8));
                return Base64.encodeBase64String(res);
            } else {
                throw new RsaException("请正确传入公钥和私钥");
            }
        } catch (Exception e) {
            throw new RsaException("加密数据失败");
        }
    }

    /**
     * 解密
     *
     * @return
     */
    public static String rsaDecrypt(String publicKey, String privateKey, String letter) {
        try {
            if (null != publicKey && null == privateKey) {
                byte[] bytes = Base64.decodeBase64(publicKey);
                X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(bytes);
                KeyFactory keyFactory = KeyFactory.getInstance(CIPHER_KEY);
                PublicKey key = keyFactory.generatePublic(x509EncodedKeySpec);
                Cipher cipher = Cipher.getInstance(CIPHER_KEY);
                cipher.init(Cipher.DECRYPT_MODE, key);
                byte[] letterBytes = Base64.decodeBase64(letter);
                byte[] res = cipher.doFinal(letterBytes);
                return new String(res, StandardCharsets.UTF_8);
            } else if (null == publicKey && null != privateKey) {
                byte[] bytes = Base64.decodeBase64(privateKey);
                PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(bytes);
                KeyFactory keyFactory = KeyFactory.getInstance(CIPHER_KEY);
                PrivateKey key = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
                Cipher cipher = Cipher.getInstance(CIPHER_KEY);
                cipher.init(Cipher.DECRYPT_MODE, key);
                byte[] letterBytes = Base64.decodeBase64(letter);
                byte[] res = cipher.doFinal(letterBytes);
                return new String(res, StandardCharsets.UTF_8);
            } else {
                throw new RsaException("请正确传入公钥和私钥");
            }
        } catch (Exception e) {
            throw new RsaException("解密数据失败");
        }
    }

    private static String codeHandle(String letter, Cipher cipher, int mode) {
        try {
            if (Cipher.ENCRYPT_MODE == mode) {
                // 加密
                byte[] res = cipher.doFinal(letter.getBytes(StandardCharsets.UTF_8));
                return Base64.encodeBase64String(res);
            } else if (Cipher.DECRYPT_MODE == mode) {
                // 解密
                byte[] letterBytes = Base64.decodeBase64(letter);
                byte[] res = cipher.doFinal(letterBytes);
                return new String(res, StandardCharsets.UTF_8);
            } else {
                throw new RsaException("不支持当前模式");
            }
        } catch (Exception e) {
            throw new RsaException("不支持当前模式");
        }
    }
}
